package src;

import java.nio.ByteBuffer;

public class AESCipherCTR {
    private final byte[] expandedSecretKey;
    private final byte[] counter;

    private byte[] encryptedText;

    public AESCipherCTR(byte[] secretKey, byte[] counter) {
        expandedSecretKey = new byte[16 * 11];
        System.arraycopy(secretKey, 0, expandedSecretKey, 0, 16);
        KeyExpansion.key_expansion(expandedSecretKey);
        this.counter = counter;
    }

    public AESCipherCTR(byte[] encryptedText,byte[] secretKey, byte[] counter) {
        expandedSecretKey = new byte[16 * 11];
        System.arraycopy(secretKey, 0, expandedSecretKey, 0, 16);
        KeyExpansion.key_expansion(expandedSecretKey);
        this.counter = counter;
    }
    private void counterIncrementing(byte[] counter) {
        for (int i = counter.length - 1; i >= 0; i--) {
            counter[i] = (byte) (counter[i] + 1);
            if (counter[i] == 0 && i != 0) {
                continue;
            } else {
                break;
            }
        }
    }

    private final AESCipher cipher = new AESCipher();
    private final byte[] encryptedCounter = new byte[16];

    private void aesCtr(byte[] data, byte[] out, int outOff) {
        cipher.encrypt(counter, expandedSecretKey, encryptedCounter);
        counterIncrementing(counter);
        for (int i = 0; i < 16; i++) {
            out[i + outOff] = (byte) (data[i] ^ encryptedCounter[i]);
        }
    }

    private byte[] buffer = new byte[16];
    private int bufferSize = 0;


    public byte[] update(byte[] data, int off, int len) {
        int leftlen = len, offset = off, outOff = 0;
        byte[] out = new byte[(bufferSize + len) / 16 * 16];
        while (leftlen > 0) {
            int copysize = 16 - bufferSize;
            if (copysize > leftlen) {
                copysize = leftlen;
            }
            System.arraycopy(data, offset, buffer, bufferSize, copysize);
            leftlen -= copysize;
            offset += copysize;
            bufferSize += copysize;
            if (bufferSize == 16) {
                aesCtr(buffer, out, outOff);
                bufferSize = 0;
                outOff += 16;
            }
        }
        return out;
    }

    public void doFinal() {
        if (bufferSize != 0) {
            throw new IllegalStateException("Buffer is not empty but " + bufferSize + " bytes long");
        }
    }

    public int getBlockSize() {
        return 16;
    }

    public void update(ByteBuffer input, ByteBuffer output) {
        var buf = new byte[getBlockSize()];
        while (input.remaining() >= getBlockSize()) {
            input.get(buf);
            output.put(update(buf, 0, getBlockSize()));
        }
    }
}
